import datetime

from sdv.multi_table import HMASynthesizer
import openpyxl

from tool_utils import *





def all_database_import(st):
    st.title("数据源管理")
    sqlite_tool = get_sql_lite_tool()

    # 获取所有数据源
    data_sources = sqlite_tool.get_table_data('database_connections', limit=-1).get('data')



    all_tables_metadata_list = sqlite_tool.get_table_data('database_all_table_metadata', limit=-1).get('data')
    has_metadata_dict = {}
    if all_tables_metadata_list:
        for table_metadata in all_tables_metadata_list:

            if table_metadata.get('sdv_metadata') is not None:
                has_metadata_dict[table_metadata['database_id']] = True


    # 显示数据源列表
    st.subheader("数据源列表")
    for source in data_sources:
        # 检查是否已创建metadata
        has_metadata = False
        try:
            # 尝试获取表元数据
            if has_metadata_dict[source['id']]:  # 如果有表说明已经创建了metadata
                has_metadata = True
        except:
            pass

        # 显示数据源信息
        with st.expander(f"{source['name']}  {'-✅ 已创建metadata' if has_metadata else ''}"):
            st.write(f"类型: {source['type']}")
            st.write(f"主机: {source['host']}")
            st.write(f"数据库: {source['database']}")

            if has_metadata:
                # get metadata
                all_table_metadata_dict = sqlite_tool.get_table_data('database_all_table_metadata',
                                                                     where_clause=f"database_id={source['id']}").get(
                    'data')[0]
                ## load
                all_table_metadata = MultiTableMetadata.load_from_dict(
                    json.loads(all_table_metadata_dict['sdv_metadata']))
                if st.button("展示结构",key=f"show_structure_{source['id']}"):
                    old_value = st.session_state.get(f"check_{source['id']}",False)
                    new_value = not old_value
                    if new_value:
                        st.write(all_table_metadata.visualize(
                            show_table_details='full',
                            show_relationship_labels=True,
                        ))
                    st.session_state[f"check_{source['id']}"] = new_value
                st.subheader("目标数据源")

                ## 获取
                type = source['type']
                target_source_list = get_data_sources_by_type(data_sources, type)
                append_csv_data_source(target_source_list)
                target_source_name = st.selectbox("选择目标数据源",
                                             [s['name'] for s in target_source_list],
                                             key=f"target_{source['id']}")
                target_source = next(s for s in target_source_list if s['name'] == target_source_name)
                st.subheader("config")
                col1,col2,col3,col4,col5 = st.columns([2,2,2,2,2])
                with col1:
                    learn_rate = st.text_input("学习比例",key=f"learn_rate_{source['id']}",value=0.5)
                with col2:
                    min_learn_num = st.text_input("最小学习数量", key=f"learn_num_min_{source['id']}", value=100)
                with col3:
                    max_learn_rate = st.text_input("最大学习数量",key=f"learn_num_max_{source['id']}",value=5000)
                with col4:
                    create_rate = st.text_input("生成比例", key=f"create_rate_{source['id']}",value=2)
                with col5:
                    language = st.multiselect("语言",
                                                     options=['zh_CN', 'en_US'],
                                                     default = 'en_US',
                                                     key=f"language_{source['id']}")

                # 添加操作按钮
                if st.button("开始处理", key=f"process_create_{source['id']}"):


                    model = HMASynthesizer(all_table_metadata,locales=language)
                    ## get data from database

                    table_data_dict = get_all_tables_data(database_id=source['id'],min_count=int(min_learn_num),max_count=int(max_learn_rate),rate=float(learn_rate))

                    model.fit(table_data_dict)
                    datas = model.sample(float(create_rate))

                    if target_source['name'] == 'CSV':
                        output = BytesIO()
                        with pd.ExcelWriter(output, engine='openpyxl') as writer:
                            for key, df1 in datas.items():
                                df1.to_excel(writer, sheet_name=key, index=False)
                        print(datas)
                        st.download_button(
                            label="下载Excel",
                            data=output.getvalue(),
                            file_name=f'gen_data_{datetime.datetime.now().strftime("%Y%m%d%H%M%S")}.xlsx',
                            mime='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
                        )

                    else:
                        target_engine = get_sql_databaseTool(target_source['id']).engine
                        write_mode = '覆盖'
                        progress_bar = st.progress(0)
                        total = len(datas)
                        idx = 0
                        ## 插入数据
                        for table_name,synthetic_data in datas.items():
                            # 连接目标数据库

                            ## has_auto_increment_pk
                            # 生成数据
                            '''
                            if has_auto_increment_pk(source_table):
                                ## delete_pk
                                synthetic_data = synthetic_data.drop(columns=source_table.primary_key.columns.keys())'''

                            # create table
                            create_table_if_not_exist(source['id'], table_name, target_source['id'], table_name)

                            if write_mode == "覆盖":
                                target_tool = get_sql_databaseTool(target_source['id'])
                                target_tool.truncate_table(table_name)
                                if target_tool.is_sqlserver():
                                    with target_tool.engine.connect() as conn:

                                        conn.execute(text(f"SET IDENTITY_INSERT {target_tool.get_full_table_name(table_name)} ON"))
                                        # 使用append模式写入（保证不重建表）
                                        synthetic_data.to_sql(
                                            name=table_name,
                                            con = conn,
                                            if_exists="append",
                                            index=False,
                                            schema=target_source['schema_name']
                                        )
                                        conn.execute(
                                            text(f"SET IDENTITY_INSERT {target_tool.get_full_table_name(table_name)} OFF"))
                                        conn.commit()
                                else:
                                    # 使用append模式写入（保证不重建表）
                                    synthetic_data.to_sql(
                                        name=table_name,
                                        con=target_engine,
                                        if_exists="append",
                                        index=False,
                                        schema=target_source['schema_name']
                                    )
                            else:
                                # 直接追加数据
                                synthetic_data.to_sql(
                                    name=table_name,
                                    con=target_engine,
                                    if_exists="append",
                                    index=False,
                                    schema=target_source['schema_name']
                                )

                            # 更新进度
                            progress = (idx + 1) / total
                            progress_bar.progress(progress)
                            st.success(f"表 [{table_name}] → [{table_name}] 生成成功！")
                            idx += 1
                        pass
            else:
                if st.button("生成metadata", key=f"process_metadata_{source['id']}"):
                    print('start.....生成metadata')
                    all_table_metadata = create_all_database_metadata(source['id'])
                    if all_table_metadata is not None:
                        data_dict = {}
                        data_dict['database_id'] = source['id']
                        data_dict['sdv_metadata'] = str(all_table_metadata)
                        if sqlite_tool.insert_data('database_all_table_metadata', data_dict):
                            st.success("metadata创建成功...")
                        else:
                            st.error("metadata创建失败...")
                    else:
                        st.error("metadata创建失败...")

            # 显示目标数据源选择

            # 示例：可以在这里添加处理逻辑
            # process_data_source(source, target_source)


